import os
from volcenginesdkarkruntime import Ark

from config import AI_KEY_HUO_SHAN
from ai.common import process_messages, process_content

from ai.consts import PLATFORM_HUO_SHAN


def get_result_by_huoshan(data, ai_info):
    os.environ["ARK_API_KEY"] = AI_KEY_HUO_SHAN
    client = Ark(
        api_key=os.environ.get("ARK_API_KEY"), 
        timeout=1800,
    )
    completion = client.chat.completions.create(
        model=ai_info['model'],
        messages=process_messages(data)
    )
    content = completion.choices[0].message.content
    return process_content(content, ai_info)
    
    
def get_huoshan_result(data, model_name):
    huoshan_ai_data = {
        "api_key": AI_KEY_HUO_SHAN,
        "model": model_name,
        "platform": PLATFORM_HUO_SHAN
    }
    return get_result_by_huoshan(data, huoshan_ai_data)

